MIRA 团队在 KDD Cup 2021 知识图谱链接预测比赛中荣获大奖！

作者 | 潘涛星

摘要

在6月16号刚结束的 KDD Cup 2021 和 OGB 官方联合举办的第一届图学习竞赛 OGB Large-Scale Challenge 中，由王杰教授指导的中科大 GraphMIRAcles 团队（蔡健宇、陈佳俊、潘涛星和张占秋同学）在知识图谱链接预测比赛中荣获大奖！代码和技术报告已公开。

比赛结果公示链接：https://ogb.stanford.edu/kddcup2021/results/#final_wikikg90m

技术报告：https://ogb.stanford.edu/paper/kddcup2021/wikikg90m_GraphMIRAcles.pdf

代码：https://github.com/MIRALab-USTC/KDDCup2021_WikiKG90M_GraphMIRAcles

比赛介绍

现实应用中存在大量以图结构形式组织的数据，而针对图结构数据的图学习方法近年也受到了学术界和工业界的高度重视。图学习算法被广泛应用到知识图谱、社交网络和推荐系统等应用中。但是这些应用中的图数据的规模正在日渐膨胀，一张图可能包含有数以十亿计的边，这给目前的图学习社区带来了巨大的挑战。为此，KDD Cup 与 OGB (Open Graph Benchmark) 团队联合举办了第一届大规模图学习比赛。

MIRA 根据团队在链接预测任务中的积累选择了本次大赛的知识图谱链接预测比赛为专攻方向，挑战大规模知识图谱对链接预测任务带来的困难。在该赛道中，主办方提供了来自真实世界的超大规模知识图谱 WikiKG90M-LSC，其中训练集就包括超过8千万个实体以及5亿个三元组。本次比赛，主办方未提供知识图谱的文本数据，而是提供实体和关系经过 RoBERTa 编码后的768维向量（下称语义向量）。

本次比赛的测试集给定三元组（头实体、关系、尾实体）中的头实体和关系，任务是对指定的1001个候选实体进行排序使得正确答案在候选实体中的排名尽可能地高。

MIRA 比赛实践

「总体方法概览」

总体框架

如上图所示，我们的方法包含三个模块——基础模型 (ComplEx-CMRC)，规则挖掘模型 (AMIE 3) 和推理模型。我们先在训练集上训练基础模型 ComplEx-CMRC，并根据在验证集上的模型表现挑选最优的模型，再利用规则挖掘算法对训练数据集进行数据增强，最后采用 ensemble 和知识蒸馏等方法构建推理模型。

「基础模型 (ComplEx-CMRC)」

基础模型

我们将基础模型设计为 Encoder-Decoder 结构，并将设计的结构命名为 ComplEx-CMRC。"CMRC" 是 Concat-MLP with Residual Connection 的缩写。为充分利用上语义信息以及图结构信息，我们将768维的语义特征与200维的图结构特征进行拼接，并且在模型中利用上残差连接的网络设计技巧，最终设计出如上图所示的 Encoder 端模型。在 Decoder 端，我们采用 ComplEx 模型。ComplEx 模型在复数域上对三元组编码的好坏进行评估。如上图所示，Encoder 端对三元组（头实体，关系，尾实体）编码分别为 ，其中 为编码向量的维度，则 ComplEx 模型按照以下公式评估三元组编码

其中 为取实部。

「规则挖掘」

为进一步提升基础模型的性能，我们采用 AIME 3 作为规则挖掘器对训练集的知识图谱进行数据增强。举个例子，

规则挖掘器可以从数据集中多次出现数据特征”夫妻总是住在同一个地方“挖掘出上述的规则。

因为数据集有5亿个三元组，由于训练机器的内存限制，难以一次性处理整个数据集，所以我们将数据集分成5张各2亿个三元组的子图，分别进行规则挖掘，并将所有挖掘到的规则合并，利用最终合并的规则构造新的三元组，进行数据增强，从而提升基础模型的性能。

「推理模型」

推理模型

如上图所示，我们的推理模型按以下三个步骤进行迭代：

• 使用上面介绍的规则挖掘方法在训练集中挖掘规则；
• 根据随机种子以及挖掘到规则置信度的差异，我们训练多个基础模型，并以对多个模型输出的取平均的方式进行 ensemble；
• 采用知识蒸馏的思路，用单模型去预测 ensemble 模型预测的结果，提升单模型的性能；

我们将整个推理模型迭代3次，并用第三次迭代得到的多个单模型进行ensemble 后作为最终的推理模型在测试集上进行预测。

我们采用了上述一整套方案后，在 KDD CUP 2021 的大规模图学习比赛中取得了好成绩，最终赢下了链接预测比赛大奖！

封面图：https://ogb.stanford.edu/assets/img/kddlsc_logo.png

作者简介：潘涛星，2019年毕业于中国科学技术大学信息科学技术学院，获得学士学位。现于中国科学技术大学电子工程与信息科学系的 MIRA Lab 实验室攻读硕士生，师从王杰教授。研究兴趣包括强化学习和大规模优化算法。

了解更多信息请浏览 MIRA Lab 主页：

https://miralab.ai/

如有合作意向或反馈建议请联系邮箱：

[e-mail]